"""
个人助理 - 记忆管理示例

这个模块实现了一个具有记忆功能的个人助理，能够：
1. 保存和检索用户偏好
2. 维护会话状态
3. 提供个性化的交互体验
"""

import json
from typing import Optional, Dict, Any
from google.adk.agents import Agent
from google.adk.tools import ToolContext


def save_user_preference(tool_context: ToolContext, preference_name: str, preference_value: str) -> str:
    """
    保存用户偏好到用户配置文件中，供将来参考。
    
    Args:
        tool_context: 工具上下文对象
        preference_name: 偏好设置的名称（如 'home_city', 'favorite_restaurant'）
        preference_value: 偏好设置的值
        
    Returns:
        确认消息
    """
    try:
        # 使用 user: 前缀的状态会在使用持久化 SessionService 时跨会话保存
        state_key = f"user:{preference_name}"
        tool_context.state[state_key] = preference_value
        
        # 同时更新一个用户偏好列表，便于管理
        preferences_key = "user:saved_preferences"
        saved_preferences = tool_context.state.get(preferences_key, [])
        if preference_name not in saved_preferences:
            saved_preferences.append(preference_name)
            tool_context.state[preferences_key] = saved_preferences
        
        return f"好的，我已经记住您的{preference_name}是{preference_value}。"
    
    except Exception as e:
        return f"抱歉，保存偏好时遇到了问题：{str(e)}"


def recall_user_info(tool_context: ToolContext, info_key: str) -> str:
    """
    从用户配置文件中回忆特定的用户信息。
    
    Args:
        tool_context: 工具上下文对象
        info_key: 要查询的信息键名
        
    Returns:
        用户信息或未找到信息的消息
    """
    try:
        # 首先从用户状态中查找
        state_key = f"user:{info_key}"
        value = tool_context.state.get(state_key)
        
        if value:
            return f"我记得您的{info_key}是：{value}"
        
        # 如果没有找到具体信息，尝试搜索记忆（如果配置了 MemoryService）
        try:
            search_results = tool_context.search_memory(query=f"用户 {info_key}")
            if search_results and hasattr(search_results, 'results') and search_results.results:
                # 处理搜索结果（简化版本）
                return f"从过去的对话中，我找到了关于您{info_key}的信息：{search_results.results[0] if search_results.results else '未找到相关信息'}"
        except Exception:
            # 如果没有配置 MemoryService 或搜索失败，继续
            pass
        
        return f"抱歉，我还没有记录您的{info_key}信息。您可以告诉我，我会记住的。"
    
    except Exception as e:
        return f"查询信息时遇到了问题：{str(e)}"


def list_user_preferences(tool_context: ToolContext) -> str:
    """
    列出已保存的所有用户偏好设置。
    
    Args:
        tool_context: 工具上下文对象
        
    Returns:
        用户偏好列表
    """
    try:
        preferences_key = "user:saved_preferences"
        saved_preferences = tool_context.state.get(preferences_key, [])
        
        if not saved_preferences:
            return "我还没有保存任何关于您的偏好信息。"
        
        preferences_list = []
        for pref_name in saved_preferences:
            state_key = f"user:{pref_name}"
            value = tool_context.state.get(state_key, "未设置")
            preferences_list.append(f"- {pref_name}: {value}")
        
        return f"以下是我记住的您的偏好设置：\n" + "\n".join(preferences_list)
    
    except Exception as e:
        return f"获取偏好列表时遇到了问题：{str(e)}"


def update_conversation_context(tool_context: ToolContext, context_key: str, context_value: str) -> str:
    """
    更新当前对话的上下文信息（仅在当前会话中有效）。
    
    Args:
        tool_context: 工具上下文对象
        context_key: 上下文键名
        context_value: 上下文值
        
    Returns:
        确认消息
    """
    try:
        # 使用无前缀的状态键，仅在当前会话中保存
        tool_context.state[context_key] = context_value
        return f"好的，我已经记录了当前对话的{context_key}：{context_value}"
    
    except Exception as e:
        return f"更新对话上下文时遇到了问题：{str(e)}"


def get_conversation_summary(tool_context: ToolContext) -> str:
    """
    获取当前对话的摘要信息。
    
    Args:
        tool_context: 工具上下文对象
        
    Returns:
        对话摘要
    """
    try:
        # 收集当前会话的非用户前缀状态
        session_data = {}
        user_data = {}
        
        for key, value in tool_context.state.items():
            if key.startswith("user:"):
                user_data[key[5:]] = value  # 移除 "user:" 前缀
            elif not key.startswith("app:") and not key.startswith("temp:"):
                session_data[key] = value
        
        summary_parts = []
        
        if user_data:
            summary_parts.append("用户信息：")
            for key, value in user_data.items():
                summary_parts.append(f"  - {key}: {value}")
        
        if session_data:
            summary_parts.append("当前对话上下文：")
            for key, value in session_data.items():
                summary_parts.append(f"  - {key}: {value}")
        
        if not summary_parts:
            return "当前没有记录的对话信息。"
        
        return "\n".join(summary_parts)
    
    except Exception as e:
        return f"获取对话摘要时遇到了问题：{str(e)}"


# 创建个人助理代理
root_agent = Agent(
    name="personal_assistant",
    model="gemini-2.0-flash",
    description="一个能够记住用户偏好并提供个性化服务的智能个人助理",
    instruction="""你是一个非常有帮助的个人助理。你具有以下能力：

1. **记忆用户偏好**：使用 'save_user_preference' 工具来记住重要的用户信息，如：
   - home_city（居住城市）
   - favorite_restaurant（喜欢的餐厅）
   - preferred_language（偏好语言）
   - work_schedule（工作时间）
   - dietary_restrictions（饮食限制）
   等等

2. **回忆用户信息**：使用 'recall_user_info' 工具来查询用户的历史偏好和信息

3. **管理对话上下文**：使用 'update_conversation_context' 工具记录当前对话的重要信息

4. **提供个性化服务**：基于记住的用户信息来个性化你的回复和建议

5. **状态总结**：使用 'get_conversation_summary' 和 'list_user_preferences' 工具来回顾信息

**使用指南**：
- 当用户分享重要的个人信息时，主动使用保存工具
- 在提供建议时，先回忆相关的用户偏好
- 根据用户的居住地、喜好等信息来个性化回复
- 保持友好、专业的态度
- 如果不确定某个信息是否重要，可以询问用户是否需要记住

记住：你的目标是成为一个真正了解用户并能提供个性化帮助的智能助理。""",
    
    tools=[
        save_user_preference,
        recall_user_info,
        list_user_preferences,
        update_conversation_context,
        get_conversation_summary
    ]
) 